package com.ruoyi.common.interceptor;

import cn.hutool.http.ContentType;
import com.alibaba.fastjson.JSON;
import com.ruoyi.common.constant.CommonConstant;
import com.ruoyi.common.utils.ip.IpUtils;
import com.ruoyi.common.vo.Result;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.BooleanUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.http.HttpHeaders;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.annotation.PostConstruct;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Collections;

@Component
@Slf4j
public class DDosInterceptor implements HandlerInterceptor {

    @Value("${ddos.enable}")
    private Boolean ddos;

    @Value("${ddos.interval_request}")
    private String intervalRequest;

    @Value("${ddos.interval_time}")
    private String intervalTime;

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    private static final String BLACK_LIST_KEY = "request_ip_blacklist";

    @PostConstruct
    public void init(){
        stringRedisTemplate.opsForValue().set("ddos", "true");
    }

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        String ddos = stringRedisTemplate.opsForValue().get("ddos");
        if(!BooleanUtils.toBoolean(ddos)){
            return true;
        }

        String ipAddr = IpUtils.getIpAddr(request);

        if (StringUtils.isNotBlank(ipAddr)){
            log.info("ipAddr:{}", ipAddr);
            if (stringRedisTemplate.opsForSet().isMember(BLACK_LIST_KEY, ipAddr)){
                writeBlackResponse(response);
                return false;
            }
        }

        if(!allowAccess(ipAddr)){
            stringRedisTemplate.opsForSet().add(BLACK_LIST_KEY, ipAddr);
            writeBlackResponse(response);
            return false;
        }

        return true;
    }

    private boolean allowAccess(String ipAddress) {
        String key = "ip:" + ipAddress;
        String script = "local current = tonumber(redis.call('incr', KEYS[1]))\n" +
                "if current == 1 then\n" +
                "    redis.call('expire', KEYS[1], ARGV[1])\n" +
                "end\n" +
                "if current > tonumber(ARGV[2]) then\n" +
                "    return 0\n" +
                "end\n" +
                "return 1";

        RedisScript<Long> redisScript = new DefaultRedisScript<>(script, Long.class);
        Long result = stringRedisTemplate.execute(redisScript, Collections.singletonList(key),
                intervalTime, intervalRequest);

        return result != null && result == 1;
    }

    private void writeBlackResponse(HttpServletResponse response) {
        try {
            String str = JSON.toJSONString(
                    new Result<>(CommonConstant.SC_BLACK, "IP黑名单，请联系管理员")
            );
            response.setHeader(HttpHeaders.CONTENT_TYPE, ContentType.JSON.getValue());
            response.getOutputStream().write(str.getBytes(StandardCharsets.UTF_8));
        } catch (IOException e) {
            log.error("writeResponse : ", e);
        }
    }


}
